# 波束形成多层感知器模型
import torch
import torch.nn as nn

class WfsMlp(nn.Module):
    def __init__(self, N:int = 8):
        '''
        N为线性阵列天线个数，输入层：天线数量、频率、间距、扫描角度（角度）
        '''
        super().__init__()
        self.l1_fc = nn.Linear(4, 16)
        self.l1_relu = nn.ReLU()
        self.l2_fc = nn.Linear(16, 64)
        self.l2_relu = nn.ReLU()
        self.l3_fc = nn.Linear(64, 256)
        self.l3_relu = nn.ReLU()
        self.l4_fc = nn.Linear(256, N)

    def forward(self, x):
        a1 = self.l1_relu(self.l1_fc(x))
        a2 = self.l2_relu(self.l2_fc(a1))
        a3 = self.l3_relu(self.l3_fc(a2))
        return self.l4_fc(a3)